## Copyright (C) 2010 Alex Opie <lx_op@orcon.net.nz>
##
## This program is free software; you can redistribute it and/or modify it
## under the terms of the GNU General Public License as published by
## the Free Software Foundation; either version 3 of the License, or (at
## your option) any later version.
##
## This program is distributed in the hope that it will be useful, but
## WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.	See the GNU
## General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with this program; see the file COPYING.  If not, see
## <http://www.gnu.org/licenses/>.

"""
Set of routines for testing iterative reconstruction techniques.
These are not necessarily going to be very fast, rather they are 
intended purely as a development aid.
"""

import numpy as np
import parallel_beam as pb

def recon (proj, theta = None, output_size = None):
	"""
	Reconstruct an image based on the projections `proj', using
	an iterative technique.
	"""
	
	if (theta == None):
		theta = np.linspace (0, 180, proj.shape [1], endpoint = False)
	
	theta *= np.pi / 180.
	
	if (output_size == None):
		output_size = pb._determine_outsize (proj, theta)
	
	estimate = np.zeros ((output_size, output_size))
	unit_image = np.ones ((output_size, output_size))
	n = proj.shape [0]
	
	for i in range (len (theta)):
		current = project (estimate, theta [i], n)
		diff = proj [:, i] - current
		scaling = project (unit, theta [i], n)
		scaled_diff = diff / scaling
		scaled_diff [np.isnan (scaled_diff)] = 0
		back_project (scaled_diff, theta [i], estimate, 'linear')
	
	return estimate_image
	
def _project (image, angle, proj_len):
	"""
	Calculates the projection of `image' at angle `angle'.
	`angle' in radians.
	"""
	n, m = np.shape (image)
	
	# Center of image
	xc = (m + 1) / 2
	yc = (n + 1) / 2
	
	# Divide each pixel into 2x2 subpixels
	img_hires = (image.repeat (2, axis = 0)).repeat (2, axis = 1)
	
	Y,X = np.mgrid [0 : (n - 0.5) : 2j * n,  0 : (m - 0.5) : 2j * m]
	Y += 0.75 - yc
	X += 0.75 - xc
	
	Y = Y.flatten ()
	X = X.flatten ()
	img_hires = img_hires.flatten ()
	
	projection = np.zeros (proj_len)
	
	# Project each pixel to vector (-sin(th),cos(th))
	s = -np.sin (angle) * X + np.cos (angle) * Y + proj_len / 2
	indices_int = s.astype (int)
	frac = s - indices_int
	
	top_ind = np.max (indices_int)
	projection [0:top_ind + 1] =  np.bincount (indices_int, img_hires * (1 - frac))
	projection [0:top_ind + 2] += np.bincount (indices_int + 1, img_hires * frac)
		
	return projection

def _back_project (projection, angle, image, interpolation = 'linear'):
	"""
	Back-projects the values in `projection' across `image'
	at angle `angle'. `angle' in radians.
	"""
	
	#diagonal = np.ceil (dim * np.sqrt (2))
	diagonal = np.ceil (np.hypot (*image.shape))
	if (len (projection) < diagonal):
		diff = int (2 * np.ceil ((diagonal - len (projection)) / 2))
		z = np.zeros (diff / 2)
		projection = np.append (z, projection)
		projection.resize (len (projection) + diff / 2)
		del z
	
	## Create the x & y values for each pixel
	dim = image.shape [0] # Square, so either dimension works.
	centre = (dim + 1) / 2  # dim is integer, so will auto-floor
	
	y,x = np.mgrid [0:dim, 0:dim]
	x -= centre
	y -= centre
	
	## s axis for projections, needed by interp1
	s = np.arange (len (projection)) - len (projection) / 2
	
	## Sum each projection's contribution
	s_dash = y * np.cos (angle) - x * np.sin (angle)
	interpolated = pb.interp1 (s_dash, s, projection, interpolation)
	np.add (image, interpolated, image)

	# (No return value)
